import argparse

import cv2
import torch

from mmdet.apis import inference_detector, init_detector


def parse_args():
    parser = argparse.ArgumentParser(description="MMDetection webcam demo")
    parser.add_argument("config", help="test config file path")
    parser.add_argument("checkpoint", help="checkpoint file")
    parser.add_argument(
        "--device", type=str, default="cuda:0", help="CPU/CUDA device option"
    )
    parser.add_argument("--camera-id", type=int, default=0, help="camera device id")
    parser.add_argument(
        "--score-thr", type=float, default=0.5, help="bbox score threshold"
    )
    args = parser.parse_args()
    return args


def main():
    args = parse_args()

    device = torch.device(args.device)

    model = init_detector(args.config, args.checkpoint, device=device)

    camera = cv2.VideoCapture(args.camera_id)

    print('Press "Esc", "q" or "Q" to exit.')
    while True:
        ret_val, img = camera.read()
        result = inference_detector(model, img)

        ch = cv2.waitKey(1)
        if ch == 27 or ch == ord("q") or ch == ord("Q"):
            break

        model.show_result(img, result, score_thr=args.score_thr, wait_time=1, show=True)


if __name__ == "__main__":
    main()
